#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
欧意（OKX）实时行情获取类
"""

import os
import json
import time
import hmac
import base64
import hashlib
import datetime
import requests
from typing import Dict, List, Union, Optional, Callable
import websockets
import asyncio
import aiohttp
import pandas as pd


class OKXMarket:
    """
    欧意（OKX）行情数据获取类
    支持REST API和WebSocket两种方式获取数据
    """

    def __init__(self, proxy=None):
        """
        初始化OKX行情类
        
        参数:
            proxy (str, optional): 代理URL，例如 "http://localhost:10809"
        """
        # API基础URL
        self.base_url = "https://www.okx.com"
        # WebSocket URL
        self.ws_url = "wss://ws.okx.com:8443/ws/v5/public"
        
        # 代理设置
        self.proxy = proxy
        self.proxies = {
            "http": proxy,
            "https": proxy
        } if proxy else None
        
        # WebSocket连接对象
        self.ws = None
        
        # 消息处理回调函数字典
        self.callbacks = {}
        
        # WebSocket连接状态
        self.connected = False
        
        # 心跳时间
        self.ping_interval = 30
        self.last_ping_time = 0
        
        # 重连次数
        self.reconnect_count = 0
        self.max_reconnect = 10
        
        # 重连等待时间（秒）
        self.reconnect_wait = 5
        
        # 设置WebSocket代理环境变量
        if proxy:
            # 从代理URL中提取主机和端口
            if proxy.startswith('http://'):
                proxy_host_port = proxy[7:]  # 移除 'http://'
            elif proxy.startswith('https://'):
                proxy_host_port = proxy[8:]  # 移除 'https://'
            else:
                proxy_host_port = proxy
                
            # 设置环境变量
            os.environ['HTTP_PROXY'] = proxy
            os.environ['HTTPS_PROXY'] = proxy
            print(f"已设置WebSocket代理: {proxy}")
    
    async def connect_websocket(self):
        """
        连接WebSocket
        """
        # 如果已经连接，先关闭
        if self.ws is not None:
            try:
                await self.ws.close()
                print("已关闭旧的WebSocket连接")
            except Exception as e:
                print(f"关闭旧WebSocket连接时出错: {e}")
        
        # 重置连接状态
        self.connected = False
        
        # 设置WebSocket连接选项
        connect_kwargs = {}
        if self.proxy:
            connect_kwargs['proxy'] = self.proxy
        
        try:
            # 连接WebSocket
            print(f"正在连接WebSocket: {self.ws_url}")
            self.ws = await websockets.connect(
                self.ws_url,
                ping_interval=None,  # 禁用自动ping，我们将手动发送ping
                **connect_kwargs
            )
            
            # 更新连接状态
            self.connected = True
            self.last_ping_time = time.time()
            self.reconnect_count = 0
            
            print("WebSocket连接成功")
            
            # 启动心跳任务
            asyncio.create_task(self._heartbeat())
            
            # 启动消息接收任务
            asyncio.create_task(self._receive_messages())
            
            return True
        except Exception as e:
            print(f"连接WebSocket失败: {e}")
            import traceback
            traceback.print_exc()
            self.ws = None
            self.connected = False
            return False
    
    async def _heartbeat(self):
        """
        发送心跳包，保持WebSocket连接
        """
        while self.connected and self.ws is not None and not self.ws.closed:
            try:
                # 每30秒发送一次心跳
                if time.time() - self.last_ping_time > self.ping_interval:
                    await self.ws.send('ping')
                    self.last_ping_time = time.time()
                    print("已发送心跳包")
                
                # 等待一段时间
                await asyncio.sleep(5)
            except Exception as e:
                print(f"发送心跳包时出错: {e}")
                # 如果发送心跳失败，尝试重连
                if self.reconnect_count < self.max_reconnect:
                    self.reconnect_count += 1
                    print(f"尝试重连 ({self.reconnect_count}/{self.max_reconnect})...")
                    await asyncio.sleep(self.reconnect_wait)
                    await self.connect_websocket()
                else:
                    print("重连次数过多，停止重连")
                    self.connected = False
                    break
    
    async def _receive_messages(self):
        """
        接收WebSocket消息
        """
        if self.ws is None:
            print("WebSocket未连接，无法接收消息")
            return
            
        try:
            async for message in self.ws:
                try:
                    # 处理心跳响应
                    if message == 'pong':
                        print("收到心跳响应")
                        continue
                    
                    # 解析消息
                    data = json.loads(message)
                    
                    # 检查是否是订阅确认消息
                    if 'event' in data and data['event'] == 'subscribe':
                        print(f"订阅确认: {data}")
                        continue
                    
                    # 检查是否是错误消息
                    if 'event' in data and data['event'] == 'error':
                        print(f"订阅错误: {data}")
                        continue
                    
                    # 获取频道和交易对
                    if 'arg' in data:
                        channel = data['arg'].get('channel', '')
                        inst_id = data['arg'].get('instId', '')
                        
                        # 查找对应的回调函数
                        callback_key = f"{channel}:{inst_id}"
                        if callback_key in self.callbacks:
                            # 调用回调函数
                            callback = self.callbacks[callback_key]
                            if asyncio.iscoroutinefunction(callback):
                                asyncio.create_task(callback(message))
                            else:
                                callback(message)
                        else:
                            print(f"未找到回调函数: {callback_key}")
                    else:
                        print(f"未知消息格式: {data}")
                        
                except json.JSONDecodeError:
                    print(f"解析消息失败: {message}")
                except Exception as e:
                    print(f"处理消息时出错: {e}")
                    import traceback
                    traceback.print_exc()
        except websockets.exceptions.ConnectionClosed as e:
            print(f"WebSocket连接已关闭: {e}")
            # 尝试重连
            if self.reconnect_count < self.max_reconnect:
                self.reconnect_count += 1
                print(f"尝试重连 ({self.reconnect_count}/{self.max_reconnect})...")
                await asyncio.sleep(self.reconnect_wait)
                await self.connect_websocket()
            else:
                print("重连次数过多，停止重连")
                self.connected = False
        except Exception as e:
            print(f"接收消息时出错: {e}")
            import traceback
            traceback.print_exc()
            # 尝试重连
            if self.reconnect_count < self.max_reconnect:
                self.reconnect_count += 1
                print(f"尝试重连 ({self.reconnect_count}/{self.max_reconnect})...")
                await asyncio.sleep(self.reconnect_wait)
                await self.connect_websocket()
            else:
                print("重连次数过多，停止重连")
                self.connected = False
    
    async def subscribe(self, args, callback):
        """
        订阅WebSocket频道
        
        参数:
            args (List[Dict]): 订阅参数列表
            callback (Callable): 回调函数
            
        返回:
            bool: 是否成功订阅
        """
        # 确保WebSocket已连接
        if self.ws is None or self.ws.closed:
            print("WebSocket未连接，尝试连接...")
            success = await self.connect_websocket()
            if not success:
                print("无法连接WebSocket，订阅失败")
                return False
        
        # 构建订阅消息
        subscribe_msg = {
            "op": "subscribe",
            "args": args
        }
        
        try:
            # 发送订阅消息
            await self.ws.send(json.dumps(subscribe_msg))
            
            # 注册回调函数
            for arg in args:
                channel = arg.get('channel', '')
                inst_id = arg.get('instId', '')
                callback_key = f"{channel}:{inst_id}"
                self.callbacks[callback_key] = callback
                print(f"已注册回调函数: {callback_key}")
            
            return True
        except Exception as e:
            print(f"订阅频道失败: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    async def unsubscribe(self, args):
        """
        取消订阅WebSocket频道
        
        参数:
            args (List[Dict]): 订阅参数列表
            
        返回:
            bool: 是否成功取消订阅
        """
        # 确保WebSocket已连接
        if self.ws is None or self.ws.closed:
            print("WebSocket未连接，无法取消订阅")
            return False
        
        # 构建取消订阅消息
        unsubscribe_msg = {
            "op": "unsubscribe",
            "args": args
        }
        
        try:
            # 发送取消订阅消息
            await self.ws.send(json.dumps(unsubscribe_msg))
            
            # 移除回调函数
            for arg in args:
                channel = arg.get('channel', '')
                inst_id = arg.get('instId', '')
                callback_key = f"{channel}:{inst_id}"
                if callback_key in self.callbacks:
                    del self.callbacks[callback_key]
                    print(f"已移除回调函数: {callback_key}")
            
            return True
        except Exception as e:
            print(f"取消订阅频道失败: {e}")
            return False
    
    async def close(self):
        """
        关闭WebSocket连接
        """
        if self.ws is not None:
            try:
                await self.ws.close()
                print("已关闭WebSocket连接")
            except Exception as e:
                print(f"关闭WebSocket连接时出错: {e}")
            finally:
                self.ws = None
                self.connected = False
                self.callbacks = {}
    
    def _get_timestamp(self) -> str:
        """获取ISO格式的时间戳"""
        return datetime.datetime.utcnow().isoformat(timespec='milliseconds') + 'Z'
    
    def _sign(self, timestamp: str, method: str, request_path: str, body: str = '') -> str:
        """
        生成签名
        
        参数:
            timestamp (str): ISO格式的时间戳
            method (str): 请求方法，如 GET, POST
            request_path (str): 请求路径
            body (str): 请求体，GET请求为空
            
        返回:
            str: Base64编码的签名
        """
        if not self.api_secret:
            raise ValueError("API Secret is required for signature generation")
            
        message = timestamp + method + request_path + (body or '')
        mac = hmac.new(
            bytes(self.api_secret, encoding='utf-8'),
            bytes(message, encoding='utf-8'),
            digestmod=hashlib.sha256
        )
        return base64.b64encode(mac.digest()).decode('utf-8')
    
    def _get_headers(self, method: str, request_path: str, body: str = '') -> Dict:
        """
        生成请求头
        
        参数:
            method (str): 请求方法，如 GET, POST
            request_path (str): 请求路径
            body (str): 请求体，GET请求为空
            
        返回:
            Dict: 请求头字典
        """
        timestamp = self._get_timestamp()
        
        headers = {
            'Content-Type': 'application/json',
            'OK-ACCESS-KEY': self.api_key,
            'OK-ACCESS-SIGN': self._sign(timestamp, method, request_path, body),
            'OK-ACCESS-TIMESTAMP': timestamp,
            'OK-ACCESS-PASSPHRASE': self.passphrase,
            # 常规模式
            'x-simulated-trading': '0'
        }
        
        return headers
    
    async def get_ticker(self, symbol):
        """
        获取最新的ticker数据
        
        参数:
            symbol (str): 交易对名称
            
        返回:
            dict: ticker数据
        """
        try:
            url = f"{self.base_url}/api/v5/market/ticker?instId={symbol}"
            
            async with aiohttp.ClientSession() as session:
                # 设置代理
                if self.proxy:
                    connector = aiohttp.TCPConnector(ssl=False)
                    session._connector = connector
                    session._connector_owner = True
                    session._proxy = self.proxy
                    
                async with session.get(url) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data
                    else:
                        print(f"获取ticker数据失败，状态码: {response.status}")
                        return None
                        
        except Exception as e:
            print(f"获取ticker数据时出错: {e}")
            return None
    
    def get_tickers(self, inst_type: str = None) -> Dict:
        """
        获取所有产品行情信息
        
        参数:
            inst_type (str, optional): 产品类型
                SPOT：币币
                SWAP：永续合约
                FUTURES：交割合约
                OPTION：期权
            
        返回:
            Dict: 行情数据
        """
        endpoint = "/api/v5/market/tickers"
        params = {}
        if inst_type:
            params['instType'] = inst_type
        
        try:
            response = requests.get(
                self.base_url + endpoint, 
                params=params,
                proxies=self.proxies,
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"获取行情列表失败: {e}")
            return {"code": "999999", "msg": f"请求异常: {str(e)}", "data": []}
    
    def get_orderbook(self, symbol: str, depth: int = 5) -> Dict:
        """
        获取产品深度
        
        参数:
            symbol (str): 产品ID，如 "BTC-USDT"
            depth (int, optional): 深度档位，最大值400，默认为5
            
        返回:
            Dict: 深度数据
        """
        endpoint = "/api/v5/market/books"
        params = {'instId': symbol, 'sz': depth}
        
        try:
            response = requests.get(
                self.base_url + endpoint, 
                params=params,
                proxies=self.proxies,
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"获取深度数据失败: {e}")
            return {"code": "999999", "msg": f"请求异常: {str(e)}", "data": []}
    
    def get_klines(self, symbol: str, bar: str = '1m', limit: int = 100) -> Dict:
        """
        获取K线数据
        
        参数:
            symbol (str): 产品ID，如 "BTC-USDT"
            bar (str, optional): K线周期，默认1m
                支持：1m/3m/5m/15m/30m/1H/2H/4H/6H/12H/1D/1W/1M/3M/6M/1Y
            limit (int, optional): 获取数量，最大值1000，默认100
            
        返回:
            Dict: K线数据
        """
        endpoint = "/api/v5/market/candles"
        params = {'instId': symbol, 'bar': bar, 'limit': limit}
        
        try:
            response = requests.get(
                self.base_url + endpoint, 
                params=params,
                proxies=self.proxies,
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"获取K线数据失败: {e}")
            return {"code": "999999", "msg": f"请求异常: {str(e)}", "data": []}
    
    def get_trades(self, symbol: str, limit: int = 100) -> Dict:
        """
        获取最近成交数据
        
        参数:
            symbol (str): 产品ID，如 "BTC-USDT"
            limit (int, optional): 获取数量，最大值500，默认100
            
        返回:
            Dict: 成交数据
        """
        endpoint = "/api/v5/market/trades"
        params = {'instId': symbol, 'limit': limit}
        
        try:
            response = requests.get(
                self.base_url + endpoint, 
                params=params,
                proxies=self.proxies,
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"获取成交数据失败: {e}")
            return {"code": "999999", "msg": f"请求异常: {str(e)}", "data": []}
    
    async def subscribe_ticker(self, symbols: List[str], callback):
        """
        订阅行情频道（WebSocket）
        
        参数:
            symbols (List[str]): 产品ID列表，如 ["BTC-USDT", "ETH-USDT"]
            callback: 回调函数，接收行情数据
        """
        args = []
        for symbol in symbols:
            args.append({"channel": "tickers", "instId": symbol})
            
        await self.subscribe(args, callback)
    
    async def subscribe_orderbook(self, symbols: List[str], callback, depth: str = "5"):
        """
        订阅深度频道（WebSocket）
        
        参数:
            symbols (List[str]): 产品ID列表，如 ["BTC-USDT", "ETH-USDT"]
            callback: 回调函数，接收深度数据
            depth (str, optional): 深度，可选"5", "400", "books", "books-l2-tbt"
        """
        channel = f"books{'-' + depth if depth not in ['5', 'books'] else ''}"
        args = []
        for symbol in symbols:
            args.append({"channel": channel, "instId": symbol})
            
        await self.subscribe(args, callback)
    
    async def subscribe_kline(self, symbols, callback, interval="1m"):
        """
        订阅K线数据
        
        参数:
            symbols (List[str]): 交易对列表
            callback (Callable): 回调函数
            interval (str): K线周期，默认1m
        
        返回:
            bool: 是否成功订阅
        """
        args = []
        for symbol in symbols:
            args.append({
                "channel": f"candle{interval}",
                "instId": symbol
            })
        
        return await self.subscribe(args, callback)
    
    async def subscribe_trades(self, symbols, callback):
        """
        订阅交易数据
        
        参数:
            symbols (List[str]): 交易对列表
            callback (Callable): 回调函数
        
        返回:
            bool: 是否成功订阅
        """
        args = []
        for symbol in symbols:
            args.append({
                "channel": "trades",
                "instId": symbol
            })
        
        return await self.subscribe(args, callback)

    async def get_candlesticks(self, symbol, interval, limit=100, start_time=None, end_time=None):
        """
        获取K线数据
        
        参数:
            symbol (str): 交易对，如 "BTC-USDT-SWAP"
            interval (str): K线周期，如 "1m", "5m", "15m", "1H", "4H", "1D"
            limit (int): 获取数量，默认100，最大1000
            start_time (int): 开始时间戳（毫秒）
            end_time (int): 结束时间戳（毫秒）
            
        返回:
            list: K线数据列表
        """
        try:
            # 构建请求参数
            params = {
                'instId': symbol,
                'bar': interval,
                'limit': str(limit)
            }
            
            # 添加开始时间
            if start_time:
                params['after'] = str(int(start_time))
            
            # 添加结束时间
            if end_time:
                params['before'] = str(int(end_time))
            
            # 发送请求
            url = f"{self.base_url}/api/v5/market/candles"
            
            async with aiohttp.ClientSession() as session:
                async with session.get(url, params=params, proxy=self.proxy) as response:
                    if response.status == 200:
                        result = await response.json()
                        if result.get('code') == '0':
                            return result.get('data', [])
                        else:
                            print(f"获取K线数据失败: {result}")
                            return []
                    else:
                        print(f"请求失败，状态码: {response.status}")
                        return []
                    
        except Exception as e:
            print(f"获取K线数据时出错: {e}")
            import traceback
            traceback.print_exc()
            return [] 